import streamlit as st
import etcd3
import json
import time
import pymysql
import pandas as pd
from confluent_kafka import Consumer, KafkaException
from datetime import datetime

class StatusManager:
    def __init__(self, etcd_client):
        self.etcd = etcd_client
    
    def get_node_status(self):
        health_data = []
        try:
            # 从ETCD获取所有健康状态
            for value, meta in self.etcd.get_prefix("/health/"):
                data = json.loads(value)
                health_data.append({
                    "实例ID": data['id'],
                    "状态": data['status'],
                    "主机": data['host'],
                    "版本": data['version'],
                    "最后上报": pd.to_datetime(data['last_update'], unit='s')
                })
            return health_data
        except Exception as e:
            st.error(f"获取消费者状态失败: {str(e)}")
            return health_data

    def get_topic_status(self):
        status = []
        try:
            for value, meta in self.etcd.get_prefix("/status/topics/"):
                topic = meta.key.decode().split('/')[-3]
                data = json.loads(value)
                data['topic'] = topic
                status.append(data)
            return status
        except Exception as e:
            st.error(f"获取Topic状态失败: {str(e)}")
            return status

class ETCDConfigManager:
    def __init__(self):
        self.etcd_client = None
        self.current_connections = {
            'etcd': {'status': 'disconnected', 'host': '192.168.119.213', 'port': 2379},
            'mysql': {'status': 'disconnected'},
            'kafka': {'status': 'disconnected'}
        }

    # ETCD连接管理
    def connect_etcd(self, host, port):
        try:
            self.etcd_client = etcd3.client(host=host, port=port)
            self.etcd_client.status()
            self.current_connections['etcd'] = {
                'status': 'connected',
                'host': host,
                'port': port
            }
            return True
        except Exception as e:
            st.error(f"ETCD连接失败: {str(e)}")
            self.current_connections['etcd']['status'] = 'error'
            return False
    
    # 全局配置管理
    def get_global_config(self, key):
        if not self.etcd_client:
            raise Exception("请先连接ETCD")
        value, _ = self.etcd_client.get(f"/config/global/{key}")
        return json.loads(value) if value else None

    def save_global_config(self, key, value):
        if not self.etcd_client:
            raise Exception("请先连接ETCD")
        self.etcd_client.put(f"/config/global/{key}", json.dumps(value))

    # Topic配置管理
    def get_topic_config(self, topic):
        if not self.etcd_client:
            raise Exception("请先连接ETCD")
        value, _ = self.etcd_client.get(f"/config/topics/{topic}")
        return json.loads(value) if value else None

    def save_topic_config(self, topic, config):
        if not self.etcd_client:
            raise Exception("请先连接ETCD")
        self.etcd_client.put(f"/config/topics/{topic}", json.dumps(config))

    def delete_topic_config(self, topic):
        if not self.etcd_client:
            raise Exception("请先连接ETCD")
        self.etcd_client.delete(f"/config/topics/{topic}")

    def list_topics(self):
        if not self.etcd_client:
            return []
        return [
            metadata.key.decode().split('/')[-1]
            for _, metadata in self.etcd_client.get_prefix("/config/topics/")
        ]

def show_status_panel(mgr):
    st.header("📊 运行状态监控")
    if not mgr.etcd:
        st.error("ETCD客户端未初始化")
        return
    
    nodes = mgr.get_node_status()
    # 按IP分组
    ip_groups = {}
    for node in nodes:
        ip = node['host']
        if ip not in ip_groups:
            ip_groups[ip] = {
                'hosts': [],
                'topics': {}
            }
        ip_groups[ip]['hosts'].append(node)

        # 获取该节点所有Topic状态
        for value, meta in mgr.get_topic_status():
            parts = meta.key.decode().split('/')
            if len(parts) < 4:
                continue
            topic = parts[3]
            status = json.loads(value)
            
            if topic not in ip_groups[ip]['topics']:
                ip_groups[ip]['topics'][topic] = []
            ip_groups[ip]['topics'][topic].append(status)                       

    # 展示存活IP列表
    with st.expander("🌐 存活节点概览", expanded=True):
        cols = st.columns(min(4, len(ip_groups)))
        for idx, ip in enumerate(sorted(ip_groups.keys())):
            with cols[idx % 4]:
                host_count = len(ip_groups[ip]['hosts'])
                topic_count = len(ip_groups[ip]['topics'])
                st.metric(
                    label=f"IP: {ip}",
                    value=f"{host_count}个节点",
                    help=f"活跃Topic: {topic_count}个"
                )

    # 展示每个IP的详细状态
    for ip in sorted(ip_groups.keys()):
        with st.expander(f"📡 {ip} 详细状态", expanded=False):
            col1, col2 = st.columns([1, 3])
            
            with col1:
                st.subheader("节点列表")
                for node in ip_groups[ip]['hosts']:
                    #超过一分钟就算是不健康
                    delta =int(time.time()) - node['last_update']
                    status_color = "#2ecc71" if delta < 60 else "#e74c3c"
                    
                    st.markdown(f"""
                    <div style="padding:10px;border-radius:5px;margin:5px;background:{status_color}20">
                        <b>{node['host']}</b> ({node['id']})<br>
                        版本: {node['version']}<br>
                        状态: <span style="color:{status_color}">{node['status']}</span><br>
                        最后上报: {datetime.fromtimestamp(node['last_update']).strftime("%m/%d %H:%M:%S")}<br>
                        延迟: {delta:.1f}s
                    </div>
                    """, unsafe_allow_html=True)
            
            with col2:
                st.subheader("Topic消费状态")
                if ip_groups[ip]['topics']:
                    topics = []
                    for topic, status_list in ip_groups[ip]['topics'].items():
                        latest = max(status_list, key=lambda x: x['last_update'])
                        topics.append({
                            "id" : id,
                            "Topic": topic,
                            "消费总数": latest['consumer_count'],
                            "最新Offset": latest['latest_offset'],
                            #"消费延迟": latest['consumer_lag'],
                            "最后更新": datetime.fromtimestamp(latest['last_update'])
                        })
                    
                    df = pd.DataFrame(topics)
                    st.dataframe(
                        df.sort_values("最后更新", ascending=False),
                        column_config={
                            "最后更新": st.column_config.DatetimeColumn(
                                format="YYYY-MM-DD HH:mm:ss"
                            ),
                            "消费延迟": st.column_config.ProgressColumn(
                                format="%d ms",
                                min_value=0,
                                max_value=1000
                            )
                        },
                        use_container_width=True,
                        hide_index=True
                    )
                else:
                    st.info("该节点暂无活跃Topic")

    # 自动刷新
    st.markdown("""
    <script>
    setTimeout(function(){
        window.location.reload();
    }, 30000);  // 30秒刷新
    </script>
    """, unsafe_allow_html=True)

    # 手动刷新按钮
    if st.button("🔄 刷新状态"):
        st.rerun()

def show_connection_panel(mgr):
    st.header("🔌 中间件连接管理")
    
    # ETCD连接
    with st.expander("📡 ETCD连接配置", expanded=True):
        col1, col2 = st.columns([3, 1])
        with col1:
            etcd_host = st.text_input(
                "ETCD地址",
                value=mgr.current_connections['etcd']['host'],
                key="etcd_host"
            )
            etcd_port = st.number_input(
                "ETCD端口",
                value=mgr.current_connections['etcd']['port'],
                key="etcd_port"
            )
        with col2:
            st.markdown("<br>", unsafe_allow_html=True)
            if st.button("🔗 连接ETCD", key="connect_etcd"):
                if mgr.connect_etcd(etcd_host, etcd_port):
                    st.success("ETCD连接成功！")
                    st.rerun()         
        display_connection_status('ETCD', mgr.current_connections['etcd']['status'])

    # MySQL连接
    if mgr.current_connections['etcd']['status'] == 'connected':
        mysql_config = mgr.get_global_config('mysql') or {}
        with st.expander("🗄 MySQL连接配置"):
            col1, col2 = st.columns([3, 1])
            with col1:
                new_mysql = {
                    'host': st.text_input("地址", value=mysql_config.get('host', 'mysql')),
                    'port': int(st.text_input("端口", value=mysql_config.get('port', 3306))),
                    'user': st.text_input("用户名", value=mysql_config.get('user', 'root')),
                    'password': st.text_input("密码", type="password", 
                                            value=mysql_config.get('password', 'root')),
                    'database': st.text_input("数据库", value=mysql_config.get('database', 'kafka_data'))
                }
            with col2:
                st.markdown("<br>", unsafe_allow_html=True)
                if st.button("💾 保存配置并连接", key="save_mysql"):
                    try:
                        mgr.save_global_config('mysql', new_mysql)
                        # 测试连接
                        conn = pymysql.connect(**new_mysql)
                        conn.close()
                        mgr.current_connections['mysql']['status'] = 'connected'
                        st.success("MySQL连接成功！")
                    except Exception as e:
                        st.error(f"连接失败: {str(e)}")
                        mgr.current_connections['mysql']['status'] = 'error'
            display_connection_status('MySQL', mgr.current_connections['mysql']['status'])

    # Kafka连接
    if mgr.current_connections['etcd']['status'] == 'connected':
        kafka_config = mgr.get_global_config('kafka')
        with st.expander("🚀 Kafka连接配置"):
            col1, col2 = st.columns([3, 1])
            with col1:
                new_kafka_config = {
                    "bootstrap.servers": st.text_input(
                        "Broker地址",
                        value=kafka_config.get('bootstrap.servers'),
                    ),
                    "group.id": st.text_input(
                        "消费组", 
                        value=kafka_config.get('group.id', 'consumer-test')
                    ),
                    "auto.offset.reset": st.text_input(
                        "消费顺序", 
                        value=kafka_config.get('auto.offset.reset', 'latest')
                    )
                } 
            with col2:
                st.markdown("<br>", unsafe_allow_html=True)
                if st.button("💾 保存配置并测试", key="save_kafka"):
                    try:
                        mgr.save_global_config('kafka', new_kafka_config)
                        consumer = Consumer(**new_kafka_config)
                        consumer.poll(1.0)
                        consumer.close()
                        mgr.current_connections['kafka']['status'] = 'connected'
                        st.success("Kafka连接成功！")
                    except Exception as e:
                        st.error(f"连接失败: {str(e)}")
                        mgr.current_connections['kafka']['status'] = 'error'
            display_connection_status('Kafka', mgr.current_connections['kafka']['status'])

def display_connection_status(name, status):
    color_map = {
        'connected': ('✅ 已连接', '#2e7d32'),
        'disconnected': ('🟡 未连接', '#f9a825'),
        'error': ('❌ 连接失败', '#c62828')
    }
    text, color = color_map.get(status, ('⚪ 未知', 'gray'))
    st.markdown(f"""
    <div style="padding: 0.5rem; border-radius: 0.5rem; background-color: {color}; color: white;">
        {name}状态: {text}
    </div>
    """, unsafe_allow_html=True)

def show_topic_management(mgr):
    st.header("📚 Topic配置管理")
    if not mgr.etcd_client:
        st.error("ETCD客户端未初始化")
        return
    
    # 创建新配置
    with st.expander("➕ 创建新Topic配置"):
        new_topic = st.text_input("Topic名称*")
        new_table = st.text_input("目标表名*")
        new_func = st.text_area(
            "映射函数*", 
            height=300,
            value="def transform(data):\n    return data"
        )
        if st.button("📌 提交配置"):
            if not all([new_topic, new_table, new_func]):
                st.error("请填写所有必填项")
                return
            mgr.save_topic_config(new_topic, {
                "table": new_table,
                "function": new_func,
                "status": "stopped",
                "created_at": int(time.time()),
                #datetime.now().isoformat(),
                "updated_at": int(time.time()),
            })
            st.success("配置创建成功！")
            st.rerun()

    # 现有配置管理
    topics = mgr.list_topics()
    if topics:
        selected_topic = st.selectbox("选择Topic", topics)
        config = mgr.get_topic_config(selected_topic)
        
        if config:
            col1, col2 = st.columns(2)
            with col1:
                if st.button(f"▶️ 启动消费 {selected_topic}"):
                    config["status"] = "running"
                    config["updated_at"] = int(time.time())
                    mgr.save_topic_config(selected_topic, config)
            with col2:
                if st.button(f"⏹ 停止消费 {selected_topic}"):
                    config["status"] = "stopped"
                    config["updated_at"] = int(time.time())
                    mgr.save_topic_config(selected_topic, config)
            
            with st.expander("⚙️ 配置详情"):
                updated_table = st.text_input("目标表名", value=config["table"])
                updated_func = st.text_area(
                    "映射函数", 
                    value=config["function"], 
                    height=400
                )
                created_time=datetime.fromtimestamp(config['created_at']).strftime("%Y-%m-%d %H:%M:%S")
                updated_time=datetime.fromtimestamp(config['updated_at']).strftime("%Y-%m-%d %H:%M:%S")
                st.markdown(f'''创建时间:{created_time},更新时间:{updated_time}''')
                if st.button("🔄 更新配置"):
                    mgr.save_topic_config(selected_topic, {
                        **config,
                        "table": updated_table,
                        "function": updated_func,
                        "updated_at": int(time.time())
                    })
                    st.success("配置已更新")
                
                if st.button("🗑️ 删除配置", type="secondary"):
                    mgr.delete_topic_config(selected_topic)
                    st.success("配置已删除")
                    st.rerun()

            with st.expander("📊 数据预览"):
                try:
                    mysql_config = mgr.get_global_config('mysql')
                    with pymysql.connect(**mysql_config) as conn:
                        df = pd.read_sql(
                            f"SELECT * FROM {config['table']} ORDER BY id DESC LIMIT 100",
                            conn
                        )
                        st.dataframe(df, use_container_width=True)
                except Exception as e:
                    st.error(f"数据加载失败: {str(e)}")

def main():
    st.set_page_config(
        page_title="Kafka管道管理系统",
        layout="wide",
        page_icon="🚀"
    )
    
    if 'config_mgr' not in st.session_state:
        st.session_state.config_mgr = ETCDConfigManager()
    
    with st.sidebar:
        st.title("导航菜单")
        menu = st.radio(
            "功能选择",
            ["🔌 连接管理", "📚 Topic配置", "📊 运行状态"],
            index=0
        )
        
        st.markdown("---")
        st.markdown("""
        ### 连接状态
        - ETCD: {etcd_status}
        - MySQL: {mysql_status}
        - Kafka: {kafka_status}
        """.format(
            etcd_status=st.session_state.config_mgr.current_connections['etcd']['status'],
            mysql_status=st.session_state.config_mgr.current_connections['mysql']['status'],
            kafka_status=st.session_state.config_mgr.current_connections['kafka']['status']
        ))
    
    if menu == "🔌 连接管理":
        show_connection_panel(st.session_state.config_mgr)
    elif menu == "📊 运行状态":
        if st.session_state.config_mgr.current_connections['etcd']['status'] != 'connected':
            st.warning("请先连接ETCD")
            return
        show_status_panel(StatusManager(st.session_state.config_mgr.etcd_client))
    else:
        if st.session_state.config_mgr.current_connections['etcd']['status'] != 'connected':
            st.warning("请先连接ETCD")
            return
        show_topic_management(st.session_state.config_mgr)
    

if __name__ == "__main__":
    main()